{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of LoraGPT2Model were not initialized from the model checkpoint at gpt2 and are newly initialized: ['lora_layer.W_a', 'lora_layer.W_b']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': tensor([[15496,    11,   703,   389,   345,    30, 50256],\n",
      "        [   40,   716,  3734,    11,  5875,   345,     0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 0],\n",
      "        [1, 1, 1, 1, 1, 1, 1]])}\n",
      "tensor(8.5647, grad_fn=<NllLossBackward0>)\n",
      "Epoch 1/1, Loss: 8.564696311950684\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from transformers import GPT2LMHeadModel, GPT2Tokenizer\n",
    "\n",
    "# 1. 加载预训练模型和分词器\n",
    "model_name = 'gpt2'\n",
    "model = GPT2LMHeadModel.from_pretrained(model_name)  # 加载预训练的GPT-2语言模型\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(model_name)  # 加载对应的GPT-2分词器\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# 2. 定义LoRA 插入层\n",
    "class LoRALayer(nn.Module):\n",
    "    def __init__(self, hidden_size, r=4):\n",
    "        super(LoRALayer, self).__init__()\n",
    "        self.r = r\n",
    "        # 初始化两个参数：W_a和W_b，随机值乘以0.01进行缩放\n",
    "        self.W_a = nn.Parameter(torch.randn(hidden_size, r) * 0.01)\n",
    "        self.W_b = nn.Parameter(torch.randn(r, hidden_size) * 0.01)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # 返回LoRA层的输出：x + (x @ W_a) @ W_b\n",
    "        return x + (x @ self.W_a) @ self.W_b\n",
    "\n",
    "# 3. 在GPT-2模型中插入LoRA层 (只对其中一个注意力模块进行适配)\n",
    "class LoraGPT2Model(GPT2LMHeadModel):\n",
    "    def __init__(self, config):\n",
    "        super(LoraGPT2Model, self).__init__(config)\n",
    "        # 根据模型配置插入LoRA层\n",
    "        self.lora_layer = LoRALayer(config.hidden_size)\n",
    "\n",
    "    def forward(self, input_ids, attention_mask=None, labels=None):\n",
    "        # 调用原GPT-2模型的transformer进行前向传播\n",
    "        outputs = self.transformer(input_ids, attention_mask=attention_mask)\n",
    "        hidden_states = outputs[0]  # 获取隐藏状态\n",
    "        \n",
    "        # 通过LoRA层\n",
    "        lora_output = self.lora_layer(hidden_states)\n",
    "        \n",
    "        # 通过语言模型的预测头得到logits\n",
    "        lm_logits = self.lm_head(lora_output)\n",
    "        \n",
    "        # 计算损失（如果提供了标签）\n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            loss_fct = nn.CrossEntropyLoss()\n",
    "            loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))\n",
    "        \n",
    "        # 返回损失和logits\n",
    "        return (loss, lm_logits)\n",
    "\n",
    "# 实例化带有LoRA层的GPT-2模型\n",
    "lora_model = LoraGPT2Model.from_pretrained(model_name)\n",
    "\n",
    "# 4. 准备训练数据 (简化示例)\n",
    "train_texts = [\"Hello, how are you?\", \"I am fine, thank you!\"]  # 示例训练文本\n",
    "# 将文本编码为模型输入格式\n",
    "train_encodings = tokenizer(train_texts, truncation=True, padding=True, return_tensors=\"pt\")\n",
    "print(train_encodings)\n",
    "input_ids = train_encodings.input_ids\n",
    "labels = input_ids.clone()  # 标签与输入相同\n",
    "\n",
    "# 5. 配置优化器和损失函数\n",
    "optimizer = optim.Adam(lora_model.parameters(), lr=5e-5)  # 使用Adam优化器，学习率为5e-5\n",
    "\n",
    "# 6. 训练循环\n",
    "num_epochs = 1  # 设置训练周期数\n",
    "lora_model.train()  # 将模型设置为训练模式\n",
    "for epoch in range(num_epochs):\n",
    "    optimizer.zero_grad()  # 清空梯度\n",
    "    loss, logits = lora_model(input_ids, labels=labels)  # 前向传播计算损失\n",
    "    print(loss)\n",
    "    loss.backward()  # 反向传播计算梯度\n",
    "    optimizer.step()  # 更新参数\n",
    "    print(f\"Epoch {epoch + 1}/{num_epochs}, Loss: {loss.item()}\")  # 打印当前周期的损失\n",
    "\n",
    "# 7. 训练完成后，保存模型\n",
    "# lora_model.save_pretrained('./lora_gpt2')  # 保存训练后的模型\n",
    "# tokenizer.save_pretrained('./lora_gpt2')  # 保存分词器\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[15496,    11,   703,   389,   345,    30, 50256],\n",
      "        [   40,   716,  3734,    11,  5875,   345,     0]])\n",
      "I am fine, thank you!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\PandaKun\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\transformers\\generation\\utils.py:1168: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello, how are you?\n",
      "\n",
      "I'm a little bit of a nerd. I'm a\n"
     ]
    }
   ],
   "source": [
    "print(input_ids)\n",
    "print(tokenizer.decode(input_ids[1]))\n",
    "\n",
    "input_text = \"Hello, how are you?\"\n",
    "inputs = tokenizer.encode(input_text, return_tensors='pt')\n",
    "\n",
    "output = model.generate(inputs,  num_return_sequences=1)\n",
    "# 解码生成的文本\n",
    "generated_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
    "print(generated_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPT2LMHeadModel(\n",
       "  (transformer): GPT2Model(\n",
       "    (wte): Embedding(50257, 768)\n",
       "    (wpe): Embedding(1024, 768)\n",
       "    (drop): Dropout(p=0.1, inplace=False)\n",
       "    (h): ModuleList(\n",
       "      (0-11): 12 x GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D(nf=2304, nx=768)\n",
       "          (c_proj): Conv1D(nf=768, nx=768)\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D(nf=3072, nx=768)\n",
       "          (c_proj): Conv1D(nf=768, nx=3072)\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
